#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
坐标转换工具 - 将其他格式的标注转换为YOLO格式
"""

import cv2
import os
import json
from pathlib import Path

def convert_voc_to_yolo(voc_file, image_width, image_height, class_mapping):
    """
    将VOC格式的XML标签转换为YOLO格式
    
    Args:
        voc_file: VOC格式XML文件路径
        image_width: 图像宽度
        image_height: 图像高度
        class_mapping: 类别名称到ID的映射字典
    """
    import xml.etree.ElementTree as ET
    
    tree = ET.parse(voc_file)
    root = tree.getroot()
    
    yolo_labels = []
    
    for obj in root.findall('object'):
        class_name = obj.find('name').text
        
        if class_name in class_mapping:
            class_id = class_mapping[class_name]
            
            bbox = obj.find('bndbox')
            xmin = int(bbox.find('xmin').text)
            ymin = int(bbox.find('ymin').text)
            xmax = int(bbox.find('xmax').text)
            ymax = int(bbox.find('ymax').text)
            
            # 转换为YOLO格式
            center_x = (xmin + xmax) / (2.0 * image_width)
            center_y = (ymin + ymax) / (2.0 * image_height)
            width = (xmax - xmin) / image_width
            height = (ymax - ymin) / image_height
            
            yolo_labels.append(f"{class_id} {center_x:.6f} {center_y:.6f} {width:.6f} {height:.6f}")
    
    return yolo_labels

def convert_coco_to_yolo(coco_annotations, image_width, image_height, class_mapping):
    """
    将COCO格式的JSON标注转换为YOLO格式
    
    Args:
        coco_annotations: COCO格式的标注列表
        image_width: 图像宽度
        image_height: 图像高度
        class_mapping: 类别名称到ID的映射字典
    """
    yolo_labels = []
    
    for ann in coco_annotations:
        class_name = ann['category_id']
        
        if class_name in class_mapping:
            class_id = class_mapping[class_name]
            
            # COCO格式：[x, y, width, height]
            x, y, w, h = ann['bbox']
            
            # 转换为YOLO格式
            center_x = (x + w / 2.0) / image_width
            center_y = (y + h / 2.0) / image_height
            width = w / image_width
            height = h / image_height
            
            yolo_labels.append(f"{class_id} {center_x:.6f} {center_y:.6f} {width:.6f} {height:.6f}")
    
    return yolo_labels

def convert_csv_to_yolo(csv_file, class_mapping):
    """
    将CSV格式的标注转换为YOLO格式
    
    CSV格式：image_name,class_name,xmin,ymin,xmax,ymax
    
    Args:
        csv_file: CSV文件路径
        class_mapping: 类别名称到ID的映射字典
    """
    import pandas as pd
    
    df = pd.read_csv(csv_file)
    
    grouped = df.groupby('image_name')
    
    results = {}
    
    for image_name, group in grouped:
        # 假设所有图像尺寸相同（需要根据实际情况调整）
        # 这里需要你自己提供图像尺寸信息
        image_width = 1920  # 默认值，需要根据实际情况修改
        image_height = 1080  # 默认值，需要根据实际情况修改
        
        yolo_labels = []
        
        for _, row in group.iterrows():
            class_name = row['class_name']
            
            if class_name in class_mapping:
                class_id = class_mapping[class_name]
                
                xmin = row['xmin']
                ymin = row['ymin']
                xmax = row['xmax']
                ymax = row['ymax']
                
                # 转换为YOLO格式
                center_x = (xmin + xmax) / (2.0 * image_width)
                center_y = (ymin + ymax) / (2.0 * image_height)
                width = (xmax - xmin) / image_width
                height = (ymax - ymin) / image_height
                
                yolo_labels.append(f"{class_id} {center_x:.6f} {center_y:.6f} {width:.6f} {height:.6f}")
        
        results[image_name] = yolo_labels
    
    return results

def batch_convert_images_to_labels(image_dir, output_dir, class_mapping, image_width=1920, image_height=1080):
    """
    批量处理图像，获取图像尺寸并生成标签文件
    
    Args:
        image_dir: 图像目录路径
        output_dir: 标签文件输出目录
        class_mapping: 类别名称到ID的映射字典
        image_width: 默认图像宽度
        image_height: 默认图像高度
    """
    image_path = Path(image_dir)
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # 支持的图像格式
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
    
    for img_file in image_path.iterdir():
        if img_file.suffix.lower() in image_extensions:
            try:
                # 获取图像尺寸
                img = cv2.imread(str(img_file))
                if img is not None:
                    h, w = img.shape[:2]
                    print(f"{img_file.name}: {w}x{h}")
                else:
                    print(f"警告: 无法读取图像 {img_file.name}，使用默认尺寸")
                    w, h = image_width, image_height
                
                # 这里应该根据实际需求添加标注转换逻辑
                # 目前只是创建空的标签文件作为示例
                label_file = output_path / f"{img_file.stem}.txt"
                with open(label_file, 'w') as f:
                    f.write("")  # 空的标签文件
                
            except Exception as e:
                print(f"处理 {img_file.name} 时出错: {e}")

def validate_yolo_labels(label_dir, image_dir):
    """
    验证YOLO标签文件的格式和内容
    
    Args:
        label_dir: 标签文件目录
        image_dir: 图像文件目录
    """
    label_path = Path(label_dir)
    image_path = Path(image_dir)
    
    errors = []
    warnings = []
    
    # 获取所有图像文件
    image_files = {}
    for ext in {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}:
        for img_file in image_path.glob(f"*{ext}"):
            image_files[img_file.stem] = img_file
    
    for label_file in label_path.glob("*.txt"):
        image_name = label_file.stem
        
        if image_name not in image_files:
            errors.append(f"标签文件 {label_file.name} 对应的图像文件不存在")
            continue
        
        try:
            # 获取图像尺寸
            img = cv2.imread(str(image_files[image_name]))
            if img is not None:
                img_height, img_width = img.shape[:2]
            else:
                errors.append(f"无法读取图像文件: {image_files[image_name]}")
                continue
            
            # 读取并验证标签
            with open(label_file, 'r') as f:
                lines = f.readlines()
            
            for line_num, line in enumerate(lines, 1):
                line = line.strip()
                if not line:
                    continue
                
                parts = line.split()
                if len(parts) != 5:
                    errors.append(f"{label_file.name}:{line_num} 格式错误，应该有5个值")
                    continue
                
                try:
                    class_id = int(parts[0])
                    center_x = float(parts[1])
                    center_y = float(parts[2])
                    width = float(parts[3])
                    height = float(parts[4])
                    
                    # 验证坐标范围
                    if not (0 <= center_x <= 1):
                        errors.append(f"{label_file.name}:{line_num} center_x 超出范围 [0,1]")
                    if not (0 <= center_y <= 1):
                        errors.append(f"{label_file.name}:{line_num} center_y 超出范围 [0,1]")
                    if not (0 <= width <= 1):
                        errors.append(f"{label_file.name}:{line_num} width 超出范围 [0,1]")
                    if not (0 <= height <= 1):
                        errors.append(f"{label_file.name}:{line_num} height 超出范围 [0,1]")
                    
                    # 转换为像素坐标进行边界检查
                    x_center_px = center_x * img_width
                    y_center_px = center_y * img_height
                    width_px = width * img_width
                    height_px = height * img_height
                    
                    if width_px <= 0 or height_px <= 0:
                        errors.append(f"{label_file.name}:{line_num} 边界框尺寸无效")
                    
                except ValueError:
                    errors.append(f"{label_file.name}:{line_num} 数值格式错误")
        
        except Exception as e:
            errors.append(f"处理标签文件 {label_file.name} 时出错: {e}")
    
    # 输出验证结果
    if errors:
        print("发现错误:")
        for error in errors:
            print(f"  ❌ {error}")
    
    if warnings:
        print("发现警告:")
        for warning in warnings:
            print(f"  ⚠️ {warning}")
    
    if not errors and not warnings:
        print("✅ 所有标签文件格式正确!")
    
    return len(errors) == 0

if __name__ == '__main__':
    # 示例：验证标签文件
    print("验证标签文件格式...")
    
    class_mapping = {
        'car': 0,
        'truck': 1,
        'bus': 2,
        'motorcycle': 3,
        'bicycle': 4
    }
    
    # 验证当前目录的标签文件
    valid = validate_yolo_labels("labels", "images")
    print(f"验证结果: {'通过' if valid else '失败'}")